package com.aos.iot.socket.common.ack;

import cn.hutool.core.util.StrUtil;
import com.aos.iot.socket.common.handler.SocketClient;
import com.aos.iot.socket.common.message.Message;
import com.aos.iot.socket.common.scheduler.CancelableScheduler;
import io.netty.util.internal.PlatformDependent;

import java.util.Map;
import java.util.concurrent.TimeUnit;

public class AckManager {

    final Map<String, AckCallback<? extends Message>> ackCallbacks = PlatformDependent.newConcurrentHashMap();

    private final CancelableScheduler scheduler;

    public AckManager(CancelableScheduler scheduler) {
        this.scheduler = scheduler;
    }

    public <T extends Message> boolean onAck(SocketClient socketTcpClient, T message) {
        String channelId = socketTcpClient.getChannelId();
        String key = message.getMsgId();
        scheduler.cancel(StrUtil.format("{}:{}", channelId, key));
        AckCallback ackCallback = removeCallback(channelId, key);
        if (ackCallback != null) {
            ackCallback.onSuccess(message);
            return true;
        }
        return false;
    }

    public void registerAck(final String key, final String channelId, AckCallback<? extends Message> callback) {
        ackCallbacks.put(StrUtil.format("{}:{}", channelId, key), callback);
        scheduleTimeout(key, channelId, callback);
    }

    private void scheduleTimeout(final String key, final String channelId, AckCallback<? extends Message> callback) {
        if (callback.getTimeout() == -1) {
            return;
        }
        scheduler.scheduleCallback(StrUtil.format("{}:{}", channelId, key), () -> {
            AckCallback<? extends Message> cb = removeCallback(channelId, key);
            if (cb != null) {
                cb.onTimeout();
            }
        }, callback.getTimeout(), TimeUnit.SECONDS);
    }

    private AckCallback<? extends Message> removeCallback(String channelId, String key) {
        return ackCallbacks.remove(StrUtil.format("{}:{}", channelId, key));
    }
}
